import psycopg2
from dp.laplace_noise import add_laplace_noise
from dp.hashing import Hasher


class Engine:
    def __init__(
        self,
        database="test",
        user="quincy",
        password="123456",
        host="localhost",
        port=5432,
        model="glove-wiki-gigaword-50",
    ):
        self.connection = psycopg2.connect(
            database=database,
            user=user,
            password=password,
            host=host,
            port=port,
        )
        self.hasher = Hasher(model=model)

        # warm up the model
        for i in range(10):
            self.hasher.hash_query("select count(*), avg(*) from foo  ")

        self.cursor = self.connection.cursor()

    def query(self, sql, sensitivity=1, epsilon=0.1, dp="rdp"):
        # dp: rdp, pdp, tdp
        self.cursor.execute(sql)
        record = self.cursor.fetchall()
        # not group by query
        assert len(record) == 1
        record = [float(i) for i in record[0]]

        # assert len(record[0]) == 1
        # print("record is ", record)
        seed = self.hasher.hash_query(sql) % 10000 if dp == "rdp" else None
        personalized = True if dp == "pdp" else False
        return add_laplace_noise(
            record, sensitivity, epsilon, seed=seed, personalized=personalized
        )

    def __exit__(self, exc_type, exc_value, traceback):
        self.cursor.close()
        self.connection.close()


if __name__ == "__main__":
    engine = Engine()
    sql = "select count(*) from t"
    print(engine.query(sql))
